import ctypes
from ctypes import c_bool, c_double, c_float, c_int, c_ubyte, c_uint64, c_void_p, POINTER, Structure
import os
import platform
import time
import cv2
import numpy as np
# The following is the mapping between key and keycode
"""
Key0 = 0
Key1 = 1
Key2 = 2
Key3 = 3
Key4 = 4
Key5 = 5
Key6 = 6
Key7 = 7
Key8 = 8
Key9 = 9
A = 10
B = 11
C = 12
D = 13
E = 14
F = 15
G = 16
H = 17
I = 18
J = 19
K = 20
L = 21
M = 22
N = 23
O = 24
P = 25
Q = 26
R = 27
S = 28
T = 29
U = 30
V = 31
W = 32
X = 33
Y = 34
Z = 35
F1 = 36
F2 = 37
F3 = 38
F4 = 39
F5 = 40
F6 = 41
F7 = 42
F8 = 43
F9 = 44
F10 = 45
F11 = 46
F12 = 47
F13 = 48
F14 = 49
F15 = 50
Down = 51
Left = 52
Right = 53
Up = 54
Apostrophe = 55
Backquote = 56
Backslash = 57
Comma = 58
Equal = 59
LeftBracket = 60
Minus = 61
Period = 62
RightBracket = 63
Semicolon = 64
Slash = 65
Backspace = 66
Delete = 67
End = 68
Enter = 69
Escape = 70
Home = 71
Insert = 72
Menu = 73
PageDown = 74
PageUp = 75
Pause = 76
Space = 77
Tab = 78
NumLock = 79
CapsLock = 80
ScrollLock = 81
LeftShift = 82
RightShift = 83
LeftCtrl = 84
RightCtrl = 85
NumPad0 = 86
NumPad1 = 87
NumPad2 = 88
NumPad3 = 89
NumPad4 = 90
NumPad5 = 91
NumPad6 = 92
NumPad7 = 93
NumPad8 = 94
NumPad9 = 95
NumPadDot = 96
NumPadSlash = 97
NumPadAsterisk = 98
NumPadMinus = 99
NumPadPlus = 100
NumPadEnter = 101
LeftAlt = 102
RightAlt = 103
LeftSuper = 104
RightSuper = 105

Unknown = 106
Used when an Unknown key has been pressed

Count = 107
"""

class Vector2(ctypes.Structure):
    _fields_ = [
        ("x", ctypes.c_double),
        ("y", ctypes.c_double),
    ]

class DinoInfo(ctypes.Structure):
    _fields_ = [
        ("position", Vector2),
        ("velocity", Vector2),
        ("collision_box_lu", Vector2),
        ("collision_box_rd", Vector2),
        ("state", ctypes.c_int),
        ("health", ctypes.c_int64),
        ("energy", ctypes.c_int64),
    ]

class FireballInfo(ctypes.Structure):
    _fields_ = [
        ("position", Vector2),
        ("velocity", Vector2),
        ("collision_box_lu", Vector2),
        ("collision_box_rd", Vector2),
    ]

class PterosaurInfo(ctypes.Structure):
    _fields_ = [
        ("position", Vector2),
        ("velocity", Vector2),
        ("collision_box_lu", Vector2),
        ("collision_box_rd", Vector2),
    ]

class CactusInfo(ctypes.Structure):
    _fields_ = [
        ("position", Vector2),
        ("velocity", Vector2),
        ("collision_box_lu", Vector2),
        ("collision_box_rd", Vector2),
    ]

class MeteoriteInfo(ctypes.Structure):
    _fields_ = [
        ("position", Vector2),
        ("velocity", Vector2),
        ("collision_box_lu", Vector2),
        ("collision_box_rd", Vector2),
    ]

class GameInfo(ctypes.Structure):
    _fields_ = [
        ("dinosaur", POINTER(DinoInfo)),
        ("fireballs", POINTER(FireballInfo)),
        ("fireball_count", c_uint64),
        ("pterosaurs", POINTER(PterosaurInfo)),   
        ("pterosaur_count", c_uint64),
        ("cactuses", POINTER(CactusInfo)),
        ("cactus_count", c_uint64),
        ("meteorites", POINTER(MeteoriteInfo)),
        ("meteorite_count", c_uint64),
        ("flag_fail", ctypes.c_bool),
        ("flag_pause", ctypes.c_bool),
        ("score", ctypes.c_uint64)
    ]


class MiniDinoGame:
    def __init__(self, to_render=True, with_window=True, to_extract_game_info=True, lib_path=None):
        """
        To run the game with no window, remember to set "to_render" to False
        If the game info is no longer needed to be extracted, set the "to_extract_game_info" to False
        """
        # Load the dylib
        system_type = platform.system()
        if lib_path is None:
            if system_type == "Windows":
                libfile_name = "mini_dino.dll"
            elif system_type == "Linux":
                libfile_name = "libmini_dino.so"
            else:
                raise RuntimeError("Unsupported os.")

        lib_path = os.path.join(os.path.dirname(__file__), libfile_name)
        self.lib = ctypes.cdll.LoadLibrary(lib_path)

        # create_game(to_render: bool, with_window: bool, to_extract_game_info: bool) -> *mut GameApp
        self.lib.create_game.argtypes = [c_bool, c_bool, c_bool]
        self.lib.create_game.restype  = c_void_p

        # initialize(&mut self)
        self.lib.initialize.argtypes = [c_void_p]
        self.lib.initialize.restype  = None

        # run(&mut self)
        self.lib.run.argtypes = [c_void_p]
        self.lib.run.restype  = None

        # step(&mut self, i_keyboard_state: *mut u8, delta_time: f64) -> bool
        self.lib.step.argtypes = [c_void_p, c_void_p, c_double]
        self.lib.step.restype  = c_bool

        # get_game_info_ptr(&self) -> *const GameInfo
        self.lib.get_game_info_ptr.argtypes = [c_void_p]
        self.lib.get_game_info_ptr.restype  = ctypes.POINTER(GameInfo)

        # get_buffer_ptr(&mut self) -> *mut u8
        self.lib.get_buffer_ptr.argtypes = [c_void_p]
        self.lib.get_buffer_ptr.restype = ctypes.POINTER(c_ubyte)

        # reset(&mut self) -> void
        self.lib.reset.argtypes = [c_void_p]
        self.lib.reset.restype  = None

        # clean_up(&mut self) -> void
        self.lib.clean_up.argtypes = [c_void_p]
        self.lib.clean_up.restype  = None

        self.game_ptr = self.lib.create_game(to_render, with_window, to_extract_game_info)
        if not self.game_ptr:
            raise RuntimeError("Failed to init_game, got null pointer!")
    
    def initialize(self):
        """Initialize the game(loading resources)"""
        self.lib.initialize(self.game_ptr)
    
    def run(self):
        """ Run the game with no block """
        self.lib.run(self.game_ptr)

    def step(self, keyboard_ptr, delta_time: float):
        """ Update the game by frame, and return if the game fails """
        return self.lib.step(self.game_ptr, keyboard_ptr, delta_time)

    def reset(self):
        """ Reset the game """
        self.lib.reset(self.game_ptr)

    def get_game_info(self):
        """ Get the game world info"""
        info_ptr = self.lib.get_game_info_ptr(self.game_ptr)
        if not info_ptr:
            return None
        
        info_val = info_ptr.contents  # 拿到结构体值

        dino_val = info_val.dinosaur.contents if info_val.dinosaur else None
        # 解析 dinosaur
        dinosaur_dict = None
        if dino_val:
            dinosaur_dict = {
                "position": (dino_val.position.x, dino_val.position.y),
                "velocity": (dino_val.velocity.x, dino_val.velocity.y),
                "collision_box_lu": (dino_val.collision_box_lu.x, dino_val.collision_box_lu.y),
                "collision_box_rd": (dino_val.collision_box_rd.x, dino_val.collision_box_rd.y),
                "state": dino_val.state,
                "health": dino_val.health,
                "energy": dino_val.energy,
            }
        
        fireballs_list = []
        if info_val.fireballs and info_val.fireball_count > 0:
            for i in range(info_val.fireball_count):
                fb = info_val.fireballs[i]
                fireballs_list.append({
                    "position": (fb.position.x, fb.position.y),
                    "velocity": (fb.velocity.x, fb.velocity.y),
                    "collision_box_lu": (fb.collision_box_lu.x, fb.collision_box_lu.y),
                    "collision_box_rd": (fb.collision_box_rd.x, fb.collision_box_rd.y),
                })

        cactuses_list = []
        if info_val.cactuses and info_val.cactus_count > 0:
            for i in range(info_val.cactus_count):
                c_ = info_val.cactuses[i]
                cactuses_list.append({
                    "position": (c_.position.x, c_.position.y),
                    "velocity": (c_.velocity.x, c_.velocity.y),
                    "collision_box_lu": (c_.collision_box_lu.x, c_.collision_box_lu.y),
                    "collision_box_rd": (c_.collision_box_rd.x, c_.collision_box_rd.y),
                })

        pterosaurs_list = []
        if info_val.pterosaurs and info_val.pterosaur_count > 0:
            for i in range(info_val.pterosaur_count):
                p_ = info_val.pterosaurs[i]
                pterosaurs_list.append({
                    "position": (p_.position.x, p_.position.y),
                    "velocity": (p_.velocity.x, p_.velocity.y),
                    "collision_box_lu": (p_.collision_box_lu.x, p_.collision_box_lu.y),
                    "collision_box_rd": (p_.collision_box_rd.x, p_.collision_box_rd.y),
                })

        meteorites_list = []
        if info_val.meteorites and info_val.meteorite_count > 0:
            for i in range(info_val.meteorite_count):
                m_ = info_val.meteorites[i]
                meteorites_list.append({
                    "position": (m_.position.x, m_.position.y),
                    "velocity": (m_.velocity.x, m_.velocity.y),
                    "collision_box_lu": (m_.collision_box_lu.x, m_.collision_box_lu.y),
                    "collision_box_rd": (m_.collision_box_rd.x, m_.collision_box_rd.y),
                })

        return {
            "dinosaur": dinosaur_dict,

            "fireballs": fireballs_list,

            "pterosaurs": pterosaurs_list,
            "cactuses": cactuses_list,
            "meteorites": meteorites_list,

            "flag_fail": info_val.flag_fail,
            "flag_pause": info_val.flag_pause,

            "score": info_val.score
        }

    def get_game_frame(self):
        """
        The returned image is stored in BGRA format,
        which is common in OpenCV
        """
        buf_ptr = self.lib.get_buffer_ptr(self.game_ptr)
        width, height = 1000, 600
        buffer = np.ctypeslib.as_array(buf_ptr, shape=(width * height * 4, ))
        image_data = buffer.reshape((height, width, 4))
        image_bgra = image_data.astype(np.uint8)
        opencv_image = cv2.cvtColor(image_bgra, cv2.COLOR_RGBA2BGRA)
        return opencv_image

    def create_keyboard_state(self, keyboard_dict):
        """
        Convert python dictionary into ctypes array,
        with a pointer to the array returned
        Example: {10: True, 19: True, 13: False}
        -> A(10), J(19) pressed, D(13) and other released
        """
        keyboard_state = (ctypes.c_ubyte * 108)()

        for key, pressed in keyboard_dict.items():
            if 0 <= key < 108:
                keyboard_state[key] = 1 if pressed else 0
        
        return ctypes.cast(keyboard_state, ctypes.POINTER(ctypes.c_ubyte))

    def close(self):
        """Close the game"""
        self.lib.clean_up(self.game_ptr)
    
    def __del__(self):
        """The GC will call it automatically"""
        self.close()
